local component = require("component")
local computer = require("computer")
local event = require("event")
local thread = require("thread")
local serialization = require("serialization")
local filesystem = require("filesystem")

-- ========== 配置区域 ==========
local SERVER_PORT = 67
local CLIENT_PORT = 68
local ARP_PORT = 69         -- ARP协议端口
local ICMP_PORT = 70        -- ICMP协议端口
local CONFIG_FILE = "/etc/network.cfg"
local DHCP_CHECK_INTERVAL = 300
local PING_TIMEOUT = 3
local ARP_TIMEOUT = 2
-- =============================

local modem = nil
local networkConfig = nil
local running = true
local dhcpThread = nil
local networkThread = nil

-- ARP表: {[IP] = {mac, timestamp}}
local arpTable = {}
local ARP_CACHE_TIME = 300 -- ARP缓存5分钟

-- 网络配置结构
local defaultConfig = {
    ip = nil,
    netmask = "255.255.255.0",
    gateway = nil,
    dns = nil,
    lease_time = 0,
    dhcp_enabled = true
}

-- ===============================
-- 工具函数
-- ===============================

-- 网卡选择功能
local function selectNetworkCard()
    local modems = {}
    for address in component.list("modem") do
        table.insert(modems, {address = address, proxy = component.proxy(address)})
    end

    if #modems == 0 then
        print("Error: No network card found!")
        return nil
    elseif #modems == 1 then
        print("Using network card: " .. modems[1].address:sub(1,8) .. "...")
        return modems[1].proxy
    else
        print("Multiple network cards detected:")
        for i, modem in ipairs(modems) do
            local isWireless = modem.isWireless and modem.isWireless() or false
            local cardType = isWireless and "(Wireless)" or "(Wired)"
            print(string.format("%d. %s %s", i, modem.address, cardType))
        end

        io.write("Enter card number (1-" .. #modems .. "): ")
        local choice = tonumber(io.read())

        if choice and choice >= 1 and choice <= #modems then
            print("Selected: " .. modems[choice].address:sub(1,8) .. "...")
            return modems[choice].proxy
        else
            print("Invalid selection, using first card")
            return modems[1].proxy
        end
    end
end

-- 加载网络配置
local function loadConfig()
    if filesystem.exists(CONFIG_FILE) then
        local file = io.open(CONFIG_FILE, "r")
        if file then
            local content = file:read("*a")
            file:close()
            local success, config = pcall(serialization.unserialize, content)
            if success and config then
                -- 合并默认配置
                for k, v in pairs(defaultConfig) do
                    if config[k] == nil then
                        config[k] = v
                    end
                end
                return config
            end
        end
    end
    return defaultConfig
end

-- 保存网络配置
local function saveConfig(config)
    local dir = filesystem.path(CONFIG_FILE)
    if not filesystem.exists(dir) then
        filesystem.makeDirectory(dir)
    end

    local file = io.open(CONFIG_FILE, "w")
    if file then
        file:write(serialization.serialize(config))
        file:close()
        return true
    end
    return false
end

-- 检查IP是否有效
local function hasValidIP()
    return networkConfig and networkConfig.ip and
            networkConfig.lease_time > computer.uptime()
end

-- IP地址工具函数
local function isInSameSubnet(ip1, ip2, netmask)
    -- 简化实现：只支持/24网络
    local subnet1 = ip1:match("(%d+%.%d+%.%d+%.)%d+")
    local subnet2 = ip2:match("(%d+%.%d+%.%d+%.)%d+")
    return subnet1 == subnet2
end

-- ===============================
-- ARP协议实现
-- ===============================

local function cleanupArpTable()
    local now = computer.uptime()
    for ip, entry in pairs(arpTable) do
        if now - entry.timestamp > ARP_CACHE_TIME then
            arpTable[ip] = nil
        end
    end
end

local function arpRequest(targetIP)
    if not networkConfig.ip then
        return nil
    end

    -- 检查ARP缓存
    cleanupArpTable()
    if arpTable[targetIP] and (computer.uptime() - arpTable[targetIP].timestamp < ARP_CACHE_TIME) then
        return arpTable[targetIP].mac
    end

    print("[ARP] Resolving " .. targetIP .. "...")

    -- 发送ARP请求
    modem.broadcast(ARP_PORT, serialization.serialize({
        type = "arp_request",
        sender_ip = networkConfig.ip,
        sender_mac = modem.address,
        target_ip = targetIP
    }))

    -- 等待ARP回复
    local deadline = computer.uptime() + ARP_TIMEOUT
    while computer.uptime() < deadline do
        local eventData = {event.pull(0.1, "modem_message")}
        if eventData[1] == "modem_message" then
            local _, _, from, port, _, payload = table.unpack(eventData)

            if port == ARP_PORT and payload then
                local success, packet = pcall(serialization.unserialize, payload)
                if success and packet and packet.type == "arp_reply" and
                        packet.target_ip == networkConfig.ip and packet.sender_ip == targetIP then

                    -- 更新ARP缓存
                    arpTable[targetIP] = {
                        mac = packet.sender_mac,
                        timestamp = computer.uptime()
                    }

                    print("[ARP] " .. targetIP .. " is at " .. packet.sender_mac:sub(1,8) .. "...")
                    return packet.sender_mac
                end
            end
        end
    end

    print("[ARP] No response from " .. targetIP)
    return nil
end

-- ===============================
-- DHCP 客户端功能
-- ===============================

local function performDHCP()
    if not networkConfig.dhcp_enabled then
        return false
    end

    print("[DHCP] Requesting IP address...")

    for attempt = 1, 3 do
        print("[DHCP] Attempt " .. attempt .. "/3")

        -- 发送DISCOVER
        modem.broadcast(SERVER_PORT, serialization.serialize({
            type = "DISCOVER",
            mac = modem.address,
            client_port = CLIENT_PORT
        }))

        -- 等待OFFER
        local deadline = computer.uptime() + 5
        while computer.uptime() < deadline do
            local eventData = {event.pull(0.1, "modem_message")}
            if eventData[1] == "modem_message" then
                local _, _, from, port, _, payload = table.unpack(eventData)

                if port == CLIENT_PORT and payload then
                    local success, packet = pcall(serialization.unserialize, payload)
                    if success and packet and packet.type == "OFFER" and packet.mac == modem.address then
                        print("[DHCP] Received offer: " .. packet.ip)

                        -- 发送REQUEST
                        modem.send(from, SERVER_PORT, serialization.serialize({
                            type = "REQUEST",
                            mac = modem.address,
                            requested_ip = packet.ip,
                            client_port = CLIENT_PORT
                        }))

                        -- 等待ACK
                        local ack_deadline = computer.uptime() + 3
                        while computer.uptime() < ack_deadline do
                            local ackData = {event.pull(0.1, "modem_message")}
                            if ackData[1] == "modem_message" then
                                local _, _, ack_from, ack_port, _, ack_payload = table.unpack(ackData)
                                if ack_port == CLIENT_PORT and ack_payload then
                                    local ack_success, ack_packet = pcall(serialization.unserialize, ack_payload)
                                    if ack_success and ack_packet and ack_packet.type == "ACK" and ack_packet.mac == modem.address then
                                        -- 更新配置
                                        networkConfig.ip = ack_packet.ip
                                        networkConfig.netmask = packet.netmask or "255.255.255.0"
                                        networkConfig.gateway = packet.gateway
                                        networkConfig.lease_time = computer.uptime() + (packet.lease_time or 3600)

                                        -- 添加网关到ARP表
                                        if networkConfig.gateway then
                                            arpTable[networkConfig.gateway] = {
                                                mac = from,  -- 假设DHCP服务器就是网关
                                                timestamp = computer.uptime()
                                            }
                                        end

                                        saveConfig(networkConfig)
                                        print("[DHCP] Success! IP: " .. ack_packet.ip)
                                        return true
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end
    end

    print("[DHCP] Failed to obtain IP address")
    return false
end

-- 后台DHCP检查线程
local function dhcpChecker()
    while running do
        if networkConfig.dhcp_enabled then
            if not hasValidIP() then
                print("[DHCP] IP lease expired or missing, requesting new lease...")
                performDHCP()
            end
        end

        -- 等待检查间隔，同时检查running状态
        local endTime = computer.uptime() + DHCP_CHECK_INTERVAL
        while computer.uptime() < endTime and running do
            os.sleep(1)
        end
    end
end

-- ===============================
-- ICMP Ping 功能
-- ===============================

local function ping(targetIP, count)
    if not networkConfig.ip then
        print("Error: No IP address configured")
        return
    end

    count = count or 4
    local sent = 0
    local received = 0
    local responses = {}

    print("PING " .. targetIP .. " from " .. networkConfig.ip)

    -- 首先进行ARP解析
    local targetMAC = nil

    if isInSameSubnet(networkConfig.ip, targetIP, networkConfig.netmask) then
        -- 同一子网，直接ARP
        targetMAC = arpRequest(targetIP)
    elseif networkConfig.gateway then
        -- 不同子网，通过网关
        print("Target is remote, using gateway " .. networkConfig.gateway)
        targetMAC = arpRequest(networkConfig.gateway)
        if targetMAC then
            print("Will send packet to gateway MAC: " .. targetMAC:sub(1,8) .. "...")
        end
    else
        print("No gateway configured for remote target")
        return
    end

    if not targetMAC then
        print("ARP failed - cannot reach " .. targetIP)
        return
    end

    for i = 1, count do
        local startTime = computer.uptime()
        local pingId = math.random(1, 65535)

        -- 发送ICMP ping到目标MAC地址
        modem.send(targetMAC, ICMP_PORT, serialization.serialize({
            type = "icmp_echo_request",
            target_ip = targetIP,
            source_ip = networkConfig.ip,
            source_mac = modem.address,
            id = pingId,
            seq = i,
            timestamp = startTime
        }))

        sent = sent + 1
        local replied = false

        -- 等待回复
        local deadline = startTime + PING_TIMEOUT
        while computer.uptime() < deadline and not replied do
            local eventData = {event.pull(0.1, "modem_message")}
            if eventData[1] == "modem_message" then
                local _, _, from, port, _, payload = table.unpack(eventData)
                if port == ICMP_PORT and payload then
                    local success, packet = pcall(serialization.unserialize, payload)
                    if success and packet and packet.type == "icmp_echo_reply" and
                            packet.id == pingId and packet.target_ip == networkConfig.ip and
                            packet.source_ip == targetIP then

                        local responseTime = (computer.uptime() - startTime) * 1000
                        table.insert(responses, responseTime)
                        print(string.format("Reply from %s: seq=%d time=%.1fms",
                                targetIP, packet.seq, responseTime))
                        received = received + 1
                        replied = true
                    end
                end
            end
        end

        if not replied then
            print("Request timeout for seq=" .. i)
        end

        if i < count then
            os.sleep(1)
        end
    end

    -- 统计信息
    print("\n--- " .. targetIP .. " ping statistics ---")
    print(string.format("%d packets transmitted, %d received, %.1f%% packet loss",
            sent, received, (sent - received) / sent * 100))

    if #responses > 0 then
        local total = 0
        local min_time = responses[1]
        local max_time = responses[1]

        for _, time in ipairs(responses) do
            total = total + time
            if time < min_time then min_time = time end
            if time > max_time then max_time = time end
        end

        local avg_time = total / #responses
        print(string.format("round-trip min/avg/max = %.1f/%.1f/%.1f ms",
                min_time, avg_time, max_time))
    end
end

-- ===============================
-- 网络协议栈后台处理
-- ===============================

local function networkListener()
    while running do
        local eventData = {event.pull(0.5, "modem_message")}

        if eventData[1] == "modem_message" and networkConfig.ip then
            local _, _, from, port, _, payload = table.unpack(eventData)

            if payload then
                local success, packet = pcall(serialization.unserialize, payload)

                if success and packet then
                    -- 处理ARP请求
                    if port == ARP_PORT and packet.type == "arp_request" and
                            packet.target_ip == networkConfig.ip then

                        print("[ARP] ARP request from " .. packet.sender_ip ..
                                " (" .. packet.sender_mac:sub(1,8) .. "...)")

                        -- 更新ARP表
                        arpTable[packet.sender_ip] = {
                            mac = packet.sender_mac,
                            timestamp = computer.uptime()
                        }

                        -- 发送ARP回复
                        modem.send(packet.sender_mac, ARP_PORT, serialization.serialize({
                            type = "arp_reply",
                            sender_ip = networkConfig.ip,
                            sender_mac = modem.address,
                            target_ip = packet.sender_ip,
                            target_mac = packet.sender_mac
                        }))

                        -- 处理ARP回复
                    elseif port == ARP_PORT and packet.type == "arp_reply" and
                            packet.target_ip == networkConfig.ip then

                        arpTable[packet.sender_ip] = {
                            mac = packet.sender_mac,
                            timestamp = computer.uptime()
                        }

                        -- 处理ICMP Echo请求
                    elseif port == ICMP_PORT and packet.type == "icmp_echo_request" and
                            packet.target_ip == networkConfig.ip then

                        print("[ICMP] Ping from " .. packet.source_ip .. " (seq=" .. (packet.seq or "?") .. ")")

                        -- 发送ICMP回复
                        modem.send(packet.source_mac, ICMP_PORT, serialization.serialize({
                            type = "icmp_echo_reply",
                            target_ip = packet.source_ip,
                            source_ip = networkConfig.ip,
                            source_mac = modem.address,
                            target_mac = packet.source_mac,
                            id = packet.id,
                            seq = packet.seq,
                            timestamp = packet.timestamp
                        }))
                    end
                end
            end
        end
    end
end

-- ===============================
-- CLI 命令处理
-- ===============================

local function showHelp()
    print("Available commands:")
    print("  ipconfig                  - Show network configuration")
    print("  ipconfig /release         - Release DHCP lease")
    print("  ipconfig /renew           - Renew DHCP lease")
    print("  ipconfig /set <ip>        - Set static IP")
    print("  ping <target> [count]     - Ping target address")
    print("  arp                       - Show ARP table")
    print("  arp <ip>                  - Resolve specific IP")
    print("  dhcp <enable|disable>     - Enable/disable DHCP")
    print("  reload                    - Reload configuration")
    print("  save                      - Save current configuration")
    print("  exit                      - Exit program")
    print("  help                      - Show this help")
end

local function showIPConfig()
    print("\nNetwork Configuration:")
    print("======================")
    print("Network Card: " .. modem.address)
    print("IP Address  : " .. (networkConfig.ip or "Not assigned"))
    print("Subnet Mask : " .. (networkConfig.netmask or "Not set"))
    print("Gateway     : " .. (networkConfig.gateway or "Not set"))
    print("DNS Server  : " .. (networkConfig.dns or "Not set"))

    if networkConfig.dhcp_enabled then
        print("DHCP        : Enabled")
        if networkConfig.lease_time > 0 then
            local remaining = networkConfig.lease_time - computer.uptime()
            if remaining > 0 then
                print("Lease Time  : " .. math.floor(remaining / 60) .. " minutes remaining")
            else
                print("Lease Time  : Expired")
            end
        end
    else
        print("DHCP        : Disabled (Static)")
    end
    print("======================\n")
end

local function showArpTable()
    cleanupArpTable()
    print("\nARP Table:")
    print("======================")
    print("IP Address      MAC Address")
    print("-------------------------------")

    local count = 0
    for ip, entry in pairs(arpTable) do
        print(string.format("%-15s %s", ip, entry.mac))
        count = count + 1
    end

    if count == 0 then
        print("No entries")
    end
    print("======================\n")
end

local function handleCommand(input)
    local parts = {}
    for word in input:gmatch("%S+") do
        table.insert(parts, word:lower())
    end

    if #parts == 0 then
        return true
    end

    local cmd = parts[1]

    if cmd == "help" then
        showHelp()

    elseif cmd == "ipconfig" then
        if #parts == 1 then
            showIPConfig()
        elseif parts[2] == "/release" then
            print("Releasing DHCP lease...")
            networkConfig.ip = nil
            networkConfig.lease_time = 0
            saveConfig(networkConfig)
            print("IP address released")
        elseif parts[2] == "/renew" then
            if networkConfig.dhcp_enabled then
                print("Renewing DHCP lease...")
                performDHCP()
            else
                print("DHCP is disabled")
            end
        elseif parts[2] == "/set" and parts[3] then
            networkConfig.ip = parts[3]
            networkConfig.dhcp_enabled = false
            networkConfig.lease_time = 0
            saveConfig(networkConfig)
            print("Static IP set to: " .. parts[3])
        else
            print("Invalid ipconfig parameters. Use: /release, /renew, or /set <ip>")
        end

    elseif cmd == "ping" then
        if parts[2] then
            local count = tonumber(parts[3]) or 4
            ping(parts[2], count)
        else
            print("Usage: ping <target> [count]")
        end

    elseif cmd == "arp" then
        if #parts == 1 then
            showArpTable()
        elseif parts[2] then
            local mac = arpRequest(parts[2])
            if mac then
                print(parts[2] .. " is at " .. mac)
            else
                print("Could not resolve " .. parts[2])
            end
        end

    elseif cmd == "dhcp" then
        if parts[2] == "enable" then
            networkConfig.dhcp_enabled = true
            saveConfig(networkConfig)
            print("DHCP enabled")
        elseif parts[2] == "disable" then
            networkConfig.dhcp_enabled = false
            saveConfig(networkConfig)
            print("DHCP disabled")
        else
            print("Usage: dhcp <enable|disable>")
        end

    elseif cmd == "reload" then
        networkConfig = loadConfig()
        print("Configuration reloaded")

    elseif cmd == "save" then
        if saveConfig(networkConfig) then
            print("Configuration saved")
        else
            print("Failed to save configuration")
        end

    elseif cmd == "exit" then
        return false

    else
        print("Unknown command: " .. cmd .. " (type 'help' for available commands)")
    end

    return true
end

-- ===============================
-- 主程序
-- ===============================

local function main()
    -- 初始化网卡
    modem = selectNetworkCard()
    if not modem then
        print("Failed to initialize network card")
        return
    end

    modem.open(CLIENT_PORT)  -- DHCP客户端端口
    modem.open(ARP_PORT)     -- ARP协议端口
    modem.open(ICMP_PORT)    -- ICMP协议端口

    -- 加载配置
    networkConfig = loadConfig()

    print("Network Client Started")
    print("Network card: " .. modem.address)
    print("Type 'help' for available commands")

    -- 启动后台线程
    networkThread = thread.create(networkListener)
    dhcpThread = thread.create(dhcpChecker)

    -- 如果没有有效IP且启用了DHCP，立即尝试获取
    if networkConfig.dhcp_enabled and not hasValidIP() then
        performDHCP()
    end

    -- CLI主循环
    while running do
        io.write("netclient> ")
        local input = io.read()

        if not input then
            break
        end

        if not handleCommand( input) then
            break
        end
    end

    -- 清理
    running = false
    if dhcpThread then dhcpThread:kill() end
    if networkThread then networkThread:kill() end
    modem.close(CLIENT_PORT)
    modem.close(ARP_PORT)
    modem.close(ICMP_PORT)
    print("Network client stopped")
end

-- 启动主程序
main()
